import numpy as np
import os
from scipy import integrate, interpolate
from Utility.PlotCurve import plot_curve
import random
import torch


def peq_cal(A_array_norm,
            cstar_rho,
            a,
            n,
            R,
            At):
    # 计算燃面面积对应的平衡压力
    p_array = np.power(
        a*cstar_rho*A_array_norm*R*R/At, 1.0/(1.0-n))
    max_p = max(p_array)
    mid_index = int(np.ceil(0.8 * len(p_array))) - 1
    mid_p = p_array[mid_index]
    inital_p = p_array[0]
    return inital_p, max_p, mid_p, p_array


def internal_ballistic_calculation(w_array_norm,
                                   Ab_array_norm,
                                   cstar_rho,
                                   r_ref,
                                   n,
                                   R,
                                   At,
                                   verify_range=True):
    # 零维内弹道计算
    p_ref = 7e6
    p_shut_off = 1.0e6
    a = r_ref/np.power(p_ref, n)
    inital_p, max_p, mid_p, p_array = peq_cal(Ab_array_norm, cstar_rho, a, n, R, At)
    if(verify_range):
        if (max_p < 2e6 or max_p > 50e6):
            #print(f"max_pressure {max_p} out of range")
            return None, None  # 如果压力顶点超出了范围，那么就不计算内弹道，直接返回None
        if (inital_p < 2e6):
            #print(f"inital_pressure {inital_p} out of range")
            return None, None  # 如果初始压力太低，那么就不计算内弹道，直接返回None
        if (mid_p < 2e6):
            return None, None  # 如果中间压力太低，那么就不计算内弹道，直接返回None
    rb_inv_array = np.array([1.0/(a*np.power(max(p, p_shut_off), n)) for p in p_array])
    t_array = integrate.cumulative_trapezoid(y=rb_inv_array,
                                             x=w_array_norm*R,
                                             initial=0)
    if(verify_range):
        if(t_array[-1]>200.0):
            return None, None # 如果燃烧时间太长，那么就不计算内弹道，直接返回None
    return t_array, p_array

def vectoring_internal_ballistic(t_array,p_array):
    #t_array和p_array是以等间距肉厚确定的，不均匀
    #将t_array和p_array转换为等时间间隔：t_vector和p_vector，共128个点
    #注意：t_vector的最大值为t_action
    t_action=t_array[-1]
    f = interpolate.interp1d(t_array, p_array, kind='slinear', bounds_error=False, fill_value=0,assume_sorted=True)
    t_vector = np.linspace(t_array[0], t_action*1.0, 128)
    p_vector = f(t_vector)
    return t_action, t_vector, p_vector

def peq_cal_pytorch(A_array_norm,
            cstar_rho,
            a,
            n,
            R,
            At):
    """
    计算燃面面积对应的平衡压力 (PyTorch版本)
    
    参数:
        A_array_norm: 归一化燃面面积 (张量)
        cstar_rho: 特征速度与密度的乘积 (张量或标量)
        a: 燃速系数 (张量或标量)
        n: 燃速压力指数 (张量或标量)
        R: 药柱半径 (张量或标量)
        At: 喷喉面积 (张量或标量)
    
    返回:
        inital_p: 初始压力 (张量)
        max_p: 最大压力 (张量)
        mid_p: 中间压力 (张量)
        p_array: 压力数组 (张量)
    """
    # 计算平衡压力 (保留梯度)
    exponent = 1.0 / (1.0 - n)
    p_array = torch.pow(a * cstar_rho * A_array_norm * R * R / At, exponent)
    
    # 计算关键压力点
    max_p = torch.max(p_array)
    mid_index = int(torch.ceil(torch.tensor(0.8 * len(p_array)))) - 1
    mid_p = p_array[mid_index]
    inital_p = p_array[0]
    
    return inital_p, max_p, mid_p, p_array


def internal_ballistic_calculation_pytorch(w_array_norm,
                                  Ab_array_norm,
                                  cstar_rho,
                                  r_ref,
                                  n,
                                  R,
                                  At,
                                  verify_range=True):
    # 零维内弹道计算
    device = w_array_norm.device  # 使用输入张量的设备
    p_ref = torch.tensor(7e6, device=device)
    p_shut_off = torch.tensor(1.0e6, device=device)
    
    # 计算a值，保留梯度
    a = r_ref / torch.pow(p_ref, n)
    
    # 调用压力平衡计算函数（需要确保peq_cal_pytorch也是PyTorch实现）
    inital_p, max_p, mid_p, p_array = peq_cal_pytorch(Ab_array_norm, cstar_rho, a, n, R, At)
    
    if(verify_range):
        # 检查压力范围条件
        if (max_p < 2e6 or max_p > 50e6):
            return None, None
        if (inital_p < 2e6):
            return None, None
        if (mid_p < 2e6):
            return None, None
    
    # 计算rb_inv_array，使用torch.where处理p_shut_off条件
    rb_inv_array = 1.0 / (a * torch.pow(torch.where(p_array > p_shut_off, p_array, p_shut_off), n))
    
    # 使用梯形法进行积分（PyTorch没有内置的cumulative_trapezoid，需要自己实现）
    dx = w_array_norm * R
    # 计算梯形面积并累加
    t_array = torch.cat([
        torch.zeros(1, device=device),
        torch.cumsum(0.5 * (rb_inv_array[:-1] + rb_inv_array[1:]) * (dx[1:] - dx[:-1]), dim=0)
    ])
    
    if(verify_range):
        if t_array[-1] > 200.0:
            return None, None
    
    return t_array, p_array

def vectoring_internal_ballistic_pytorch(t_array, p_array, t_max=None,vector_length=128):
    """
    将非均匀时间间隔的压力曲线转换为等时间间隔的向量化表示 (PyTorch版本)
    
    参数:
        t_array: 非均匀时间点 (张量)
        p_array: 对应时间点的压力值 (张量)
    
    返回:
        t_action: 总作用时间 (张量)
        t_vector: 等间隔时间向量 (128点) (张量)
        p_vector: 插值后的压力向量 (张量)
    """
    device = t_array.device
    # 计算总作用时间 (保留梯度)
    if t_max is not None:
        t_action=t_max
    else:
        t_action = t_array[-1]
    
    # 创建等间隔时间向量 (0到t_action, 128点)
    t_vector = torch.linspace(t_array[0].item(), t_action.item(), vector_length, device=device)
    
    # 线性插值实现
    # 找到每个t_vector对应的区间
    indices = torch.searchsorted(t_array, t_vector) - 1
    indices = torch.clamp(indices, 0, len(t_array)-2)
    
    # 计算插值权重
    t_lower = t_array[indices]
    t_upper = t_array[indices+1]
    alpha = (t_vector - t_lower) / (t_upper - t_lower)
    alpha = torch.clamp(alpha, 0.0, 1.0)  # 确保在边界内
    
    # 执行线性插值
    p_lower = p_array[indices]
    p_upper = p_array[indices+1]
    p_vector = p_lower + alpha * (p_upper - p_lower)
    
    return t_action, t_vector, p_vector


def generate_random(min, max):
    return random.uniform(min, max)


def generate_random_log(min, max):
    return np.exp(random.uniform(np.log(min), np.log(max)))


if __name__ == '__main__':
    from DEAD.AutoDecoder.Config import *

    grain_index = 134

    # Get all files in the directory
    files = os.listdir(f"{decode_result_save_path}")
    # 过滤出 .npz 结尾的文件
    npz_files = [file for file in files if file.endswith('.npz')]
    npz_file = [file for file in npz_files if file.startswith(
        f'{grain_index}#')][0]
    npz_file = np.load(f"{decode_result_save_path}/{npz_file}")
    w_array_norm = npz_file['burned_web_array']
    Ab_array_norm = npz_file['burning_surface_area_array']
    Ap_array_norm = npz_file['exit_port_area_array']

    for i in range(50):
        R = generate_random_log(R_min, R_max)
        J = generate_random(J_min, J_max)  # 喉通比
        At = Ap_array_norm[0]*R*R*J
        t_array, p_array = internal_ballistic_calculation(
            w_array_norm=w_array_norm,
            Ab_array_norm=Ab_array_norm,
            cstar_rho=generate_random(cstar_min*rho_min, cstar_max*rho_max),
            r_ref=generate_random(r_ref_min, r_ref_max),
            n=generate_random(n_min, n_max),
            R=R,
            At=At)

        if t_array is not None and p_array is not None:
            t_action, t_vector, p_vector = vectoring_internal_ballistic(t_array,p_array)
            plot_curve(title='Chamber Pressure vs Burned time',
                       xlabel='t /s',
                       ylabel='p /MPa',
                       x=t_vector,
                       y=p_vector*1e-6)